"""
@author : Hyunwoong
@when : 2019-10-29
@homepage : https://github.com/gusdnd852
"""
from conf import *
from util.data_loader import DataLoader
from util.tokenizer import Tokenizer
import os
import pickle

tokenizer = Tokenizer()
loader = DataLoader(ext=('.en', '.de'),
                    tokenize_en=tokenizer.tokenize_en,
                    tokenize_de=tokenizer.tokenize_de,
                    init_token='<sos>',
                    eos_token='<eos>')

train, valid, test = loader.make_dataset()
loader.build_vocab(train_data=train, min_freq=2)
train_iter, valid_iter, test_iter = loader.make_iter(train, valid, test,
                                                     batch_size=batch_size,
                                                     device=device)

src_pad_idx = loader.source.vocab.stoi['<pad>']
trg_pad_idx = loader.target.vocab.stoi['<pad>']
trg_sos_idx = loader.target.vocab.stoi['<sos>']
trg_eos_idx = loader.target.vocab.stoi['<eos>']

enc_voc_size = len(loader.source.vocab)
dec_voc_size = len(loader.target.vocab)

def save_loader(data_dir):
    obj = {
        'loader': loader,
        'src_pad_idx': src_pad_idx,
        'trg_pad_idx': trg_pad_idx,
        'trg_sos_idx': trg_sos_idx,
        'enc_voc_size': enc_voc_size,
        'dec_voc_size': dec_voc_size
    }
    print('save_loader', obj)
    pickle.dump(obj, open(os.path.join(data_dir, 'loader.pkl'), 'wb'), protocol=pickle.HIGHEST_PROTOCOL)


def load_loader(data_dir):
    loader_path = os.path.join(data_dir, 'loader.pkl')
    if not os.path.exists(loader_path):
        return
    # 临时替换会直接结束程序，没有任何提示
    # import pathlib
    # temp = pathlib.PosixPath
    # pathlib.PosixPath = pathlib.WindowsPath
    obj = pickle.load(open(loader_path, 'rb'))
    # pathlib.PosixPath = temp

    global loader, src_pad_idx, trg_pad_idx, trg_sos_idx, enc_voc_size, dec_voc_size
    loader = obj['loader']
    src_pad_idx = obj['src_pad_idx']
    trg_pad_idx = obj['trg_pad_idx']
    trg_sos_idx = obj['trg_sos_idx']
    enc_voc_size = obj['enc_voc_size']
    dec_voc_size = obj['dec_voc_size']
